import json
from flask import Flask, request, jsonify, abort
import traceback
import time
from utils.log import MyLog
from torch2trt.trt_engine import TrtModel

app = Flask(__name__)
app.config['JSON_AS_ASCII'] = False
log = MyLog('test.log', __name__).getlog()
path = '/data0/jianyu10/prompt-bert/UseModel/prompt.trt'
model_path = '/data0/jianyu10/PTM/huggingface_model_cache/chinese-roberta-wwm-ext'
trtengine = TrtModel(path, model_path)


def inference(title):
    stime = time.time()
    # response
    vec = []
    time_cost = 0
    # status
    code = 1000
    message = "success"
    # empty check
    try:
        if not title:
            code = 2001
            message = "title empty!"
        else:
            # search
            try:
                vec = trtengine(title)
            except:
                log.info('except:\n%s' % traceback.format_exc())
                code = 2002
                message = "inference error, crash!"
    except:
        code = 2001
        message = "parameters error!"
    etime = time.time()
    time_cost = (etime - stime) * 1000  # ms
    # response
    response = {"code": code, "message": message, "data": vec, "time_cost": int(time_cost)}
    return response


@app.route('/title2vec/get', methods=['GET'])
def get():
    title = request.args.get('title')
    response = inference(title)
    return json.dumps(response)


if __name__ == '__main__':
    app.run(host="0.0.0.0", port=5000)
